import os
import csv
from tqdm import tqdm
import time

project_dir = os.path.dirname(os.path.abspath(__file__))
project_dir = project_dir.replace('scripts','')
import sys
sys.path.append(project_dir)

import llmdet
t = time.time()
llmdet.load_probability(project_dir)
print('loading time:', time.time() - t)

filename = os.path.join(project_dir, 'data/train_prob.csv')

def write_to_csv(data):
    with open(filename, 'a', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(data)

score_header = ["label","Human_write", "GPT-2", "OPT", "UniLM", "LLaMA", "BART", "T5", "Bloom", "GPT-neo"]
write_to_csv([score_header])
cur_texts = []
cur_labels = []
start_row = 0

if os.path.exists(filename):
    with open(filename,'r',encoding='utf-8') as f:
        lines = csv.reader(f)
        next(lines)
        for line in lines: start_row += 1

with open(os.path.join(project_dir,'data/train_drcat_01.csv'),'r',encoding='utf-8') as f:
    lines = csv.reader(f)
    header = next(lines)
    print(header)
    
    for i,line in tqdm(enumerate(lines)):
        if i < start_row: continue
        
        cur_texts.append(line[0])
        cur_labels.append(int(line[1]))
        if len(cur_texts) >= 100:
            
            result = llmdet.detect(cur_texts, hf_model_path_prefix=f'{project_dir}/llm',
                        model_prefix=f'{project_dir}/model')
            
            # print('label:',line[1],'predict:',result)
            rows = []
            for _label, _result in zip(cur_labels,result):
                row = [_label] + [_result[k] for k in score_header[1:]]
                rows.append(row)
            write_to_csv(rows)
            cur_texts.clear()
            cur_labels.clear()
            
